
package com.simple.mqtt.server.store;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.simple.mqtt.common.tool.TopicTool;
import com.simple.mqtt.server.domain.Message;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;

/**
 * In-memory MQTT message store implementation
 * Uses ConcurrentHashMap and Guava Cache for efficient memory message management
 * Supports permanent retained messages and time-limited retained messages
 *
 * @author L.cm
 */
public class InMemoryMqttMessageStore implements IMqttMessageStore {

    /**
     * Map for storing client will messages
     * key: clientId, value: corresponding will message
     */
    private final ConcurrentMap<String, Message> willStore = new ConcurrentHashMap<>();

    /**
     * Map for storing permanent retained messages
     * key: topic, value: corresponding retained message
     */
    private final ConcurrentMap<String, Message> retainStore = new ConcurrentHashMap<>();

    /**
     * Time-limited retained message cache using Guava Cache
     * Automatically expires and cleans up based on set time
     */
    private final LoadingCache<String, Message> timedRetainStore;

    /**
     * Constructor, initialize memory storage components
     */
    public InMemoryMqttMessageStore() {
        // Initialize timed cleanup cache, check every second, max cache 2 hours
        this.timedRetainStore = CacheBuilder.newBuilder()
                .concurrencyLevel(Runtime.getRuntime().availableProcessors())
                .expireAfterWrite(2, TimeUnit.HOURS)
                .refreshAfterWrite(1, TimeUnit.SECONDS)
                .build(new CacheLoader<String, Message>() {
                    @Override
                    public Message load(String key) throws Exception {
                        return null;
                    }
                });
    }

    @Override
    public boolean addWillMessage(String clientId, Message message) {
        willStore.put(clientId, message);
        return true;
    }

    @Override
    public boolean clearWillMessage(String clientId) {
        willStore.remove(clientId);
        return true;
    }

    @Override
    public Message getWillMessage(String clientId) {
        return willStore.get(clientId);
    }

    @Override
    public boolean addRetainMessage(String topic, int timeout, Message message) {
        if (timeout <= 0) {
            retainStore.put(topic, message);
        } else {
            // Guava 不支持单个 entry 设置不同过期时间，因此我们手动实现：
            // 使用 CacheBuilder.expireAfterWrite + 手动 put + 定时刷新来模拟
            timedRetainStore.put(topic, message);
        }
        return true;
    }

    @Override
    public boolean clearRetainMessage(String topic) {
        retainStore.remove(topic);
        timedRetainStore.invalidate(topic);
        return true;
    }

    @Override
    public List<Message> getRetainMessage(String topicFilter) {
        List<Message> retainMessageList = new ArrayList<>();

        // 从永久保留消息存储中查找匹配的消息
        for (Map.Entry<String, Message> entry : retainStore.entrySet()) {
            if (TopicTool.match(topicFilter, entry.getKey())) {
                retainMessageList.add(entry.getValue());
            }
        }

        // 从带过期时间的保留消息存储中查找匹配的消息
        for (Map.Entry<String, Message> entry : timedRetainStore.asMap().entrySet()) {
            if (TopicTool.match(topicFilter, entry.getKey())) {
                retainMessageList.add(entry.getValue());
            }
        }

        return retainMessageList;
    }

    @Override
    public void clean() throws IOException {
        willStore.clear();
        retainStore.clear();
        timedRetainStore.invalidateAll();
        timedRetainStore.cleanUp(); // 显式触发清理
    }
}
